from langchain.chains import RetrievalQA
from langchain.llms import Ollama
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainFilter
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings

# 加载向量数据库
embeddings = HuggingFaceEmbeddings(model_name="bge-m3")
db = Chroma(persist_directory=".", embedding_function=embeddings)

# 创建检索器
retriever = db.as_retriever()

# 使用本地的 qllama/bge-reranker-v2-m3 模型进行重排序
reranker_llm = Ollama(model="qllama/bge-reranker-v2-m3")

# 定义一个简单的提示模板用于重排序
prompt = PromptTemplate(
    input_variables=["question", "context"],
    template="以下是与问题 '{question}' 相关的一些上下文信息：{context}。请根据相关性对这些信息进行筛选。"
)

# 创建 LLMChain 用于重排序
llm_chain = LLMChain(llm=reranker_llm, prompt=prompt)

# 使用 LLMChainFilter 进行重排序
compressor = LLMChainFilter(llm_chain=llm_chain)
compression_retriever = ContextualCompressionRetriever(base_retriever=retriever, compressor=compressor)

# 使用 Gemma4b 模型
llm = Ollama(model="gemma4b")

# 创建 RAG 链
qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=compression_retriever)

# 定义一个测试问题
query = "孙悟空有哪些本领？"
result = qa.run(query)

print(result)
